{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-0BQWhvAP2jb"
      },
      "source": [
        "\n",
        "\u003ca href=\"https://colab.research.google.com/github/google-research/t5x/blob/main/t5x/notebooks/training.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bqZYp90PIa1t"
      },
      "source": [
        "# Overview\n",
        "\n",
        "This is the second Colab in a [series of tutorials on how to use T5X](https://github.com/google-research/t5x/blob/main/docs/tutorials.md). We assume that you have already completed https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.mdx-colab-intro, or have a basic understanding of the T5X models, checkpoints, partitioner, and `InteractiveModel`.\n",
        "\n",
        "In the [previous Colab](https://colab.research.google.com/github/google-research/t5x/blob/main/t5x/notebooks/introduction.ipynb) in this tutorial series, we presented a quick and easy way to use the `InteractiveModel` to run training on natural text inputs in only a few lines of code. In this Colab, we will dive into how the `InteractiveModel` restores models from checkpoints and runs training, while also getting an introduction to the T5X trainer. It should be noted that the code snippets below exactly replicate the InteractiveModel `__init__()` and `train_step()` methods (see [source code](https://github.com/google-research/t5x/blob/main/t5x/interactive_model.py)); we expose this functionality here in order to demonstrate how various components of the T5X codebase work together to train a model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nZJbWZcfkyxI"
      },
      "source": [
        "# Set-Up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rn4C8OdjVn6e"
      },
      "source": [
        "Note: If you are a using public colab, please use its `Connect to a local runtime` option by following the [setup guide](https://github.com/google-research/t5x/blob/main/t5x/notebooks/README.md)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jIGSIHzD7YPO"
      },
      "outputs": [],
      "source": [
        "from collections.abc import Sequence\n",
        "import enum\n",
        "import functools\n",
        "import inspect\n",
        "import itertools\n",
        "import logging\n",
        "import os\n",
        "import re\n",
        "from typing import Any, Callable, Iterator, Optional, Tuple, Union\n",
        "\n",
        "import jax\n",
        "from jax import random\n",
        "from jax.experimental import multihost_utils\n",
        "import numpy as np\n",
        "import seqio\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import t5.data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xAIF2kbAVw-0"
      },
      "outputs": [],
      "source": [
        "import clu.data\n",
        "from t5x.examples.t5 import network\n",
        "import t5x\n",
        "from t5x import models\n",
        "from t5x import partitioning\n",
        "from t5x import trainer as trainer_lib\n",
        "from t5x import utils\n",
        "from t5x.infer import _extract_tokens_and_aux_values\n",
        "from t5x.infer import _Inferences\n",
        "from t5x.interactive_model import InteractiveModel\n",
        "from t5x.interactive_model import get_batches_from_seqio\n",
        "from t5x.interactive_model import get_dataset_from_natural_text_examples\n",
        "from t5x.interactive_model import get_gin_config_from_interactive_model\n",
        "from t5x.interactive_model import T5XScriptType"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S5Lb-Z1fkF5a"
      },
      "source": [
        "Before we begin, let's initialize instances of the constructor arguments for the `InteractiveModel`. As mentioned previously, this will enable us to dive into how the `InteractiveModel` restores models from checkpoints and runs training.\n",
        "\n",
        "If you don't understand the lines of code below, or have questions about how to initialize these parameters, please see the [first Colab](https://colab.research.google.com/github/google-research/t5x/blob/main/t5x/notebooks/introduction.ipynb) in this tutorial series."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ne8U8qoWkX_r"
      },
      "outputs": [],
      "source": [
        "# Define a model. The configuration below corresponds to the T5 1.1 Small model.\n",
        "t5_config = network.T5Config(\n",
        "    vocab_size=32128,\n",
        "    dtype='bfloat16',\n",
        "    emb_dim=512,\n",
        "    num_heads=6,\n",
        "    num_encoder_layers=8,\n",
        "    num_decoder_layers=8,\n",
        "    head_dim=64,\n",
        "    mlp_dim=1024,\n",
        "    mlp_activations=('gelu', 'linear'),\n",
        "    dropout_rate=0.0,\n",
        "    logits_via_embedding=False)\n",
        "module = network.Transformer(config=t5_config)\n",
        "model = t5x.models.EncoderDecoderModel(\n",
        "    module=module,\n",
        "    input_vocabulary=t5.data.get_default_vocabulary(),\n",
        "    output_vocabulary=t5.data.get_default_vocabulary(),\n",
        "    optimizer_def=t5x.adafactor.Adafactor(decay_rate=0.8, step_offset=0))\n",
        "# Define checkpoint arguments.\n",
        "checkpoint_path='gs://t5-data/pretrained_models/cbqa/small_ssm_nq/model.ckpt-1110000'\n",
        "dtype='bfloat16'\n",
        "restore_mode='specific'\n",
        "# Define a partitioner.\n",
        "partitioner=partitioning.PjitPartitioner(num_partitions=2)\n",
        "# Define additional, miscellaneous constructor arguments.\n",
        "batch_size=8\n",
        "task_feature_lengths = {'inputs': 38, 'targets': 18}\n",
        "output_dir='/tmp/output_dir'\n",
        "input_shapes = {\n",
        "    'encoder_input_tokens': np.array([8, 38]),\n",
        "    'decoder_target_tokens': np.array([8, 18]),\n",
        "    'decoder_input_tokens': np.array([8, 18]),\n",
        "    'decoder_loss_weights': np.array([8, 18])\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ib9aOi2xaCKQ"
      },
      "source": [
        "# Training Deep Dive"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-trKAuOWaGoB"
      },
      "source": [
        "Let's start by going over what happens when we initialize the InteractiveModel.\n",
        "\n",
        "The `InteractiveModel` `__init__()` method performs six main actions:\n",
        "\n",
        "\n",
        "1.   Configure and possibly create an output directory.\n",
        "2.   Initialize RNGs.\n",
        "3.   Validate the partitioner.\n",
        "4.   Create a checkpoint manager.\n",
        "5.   Restore the model from a checkpoint or initialize from scratch.\n",
        "6.   Create a trainer.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zMPQQxR6lW78"
      },
      "source": [
        "**Configuring the Output Directory** \\\n",
        "There is minimal work required to configure the output directory for our model: we simply remove double-slashes in the directory path to avoid inconsistencies and create the directory if it doesn't already exist."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3wyAW8drm1JV"
      },
      "outputs": [],
      "source": [
        "output_dir = re.sub(r\"(?\u003c!gs:)([\\/]{2,})\", \"/\", output_dir)\n",
        "if not os.path.exists(output_dir):\n",
        "  os.mkdir(output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5wGHXxSynEs4"
      },
      "source": [
        "**Initializing RNGs** \\\n",
        "Initializing RNGs is made fairly straightforward with the use of JAX random operations.\n",
        "\n",
        "\n",
        "\n",
        "We first set an initial seed using the `multihost_utils` tools, then define an RNG using the JAX `PRNGKey` utils, and finally split this RNG into two values: one each for initializing the model and training the model. This ensures that we never reuse an RNG key."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wsOqpNXUnl81"
      },
      "outputs": [],
      "source": [
        "init_random_seed = 42\n",
        "random_seed = multihost_utils.broadcast_one_to_all(np.int32(init_random_seed))\n",
        "utils.set_hardware_rng_ops()\n",
        "rng = random.PRNGKey(random_seed)\n",
        "init_rng, trainer_rng = random.split(rng, 2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uiw55rQ_n-KP"
      },
      "source": [
        "**Validating the Partitioner** \\\n",
        "\n",
        "Because we've already constructed the partitioner, we simply need to validate that it was constructed properly. In particular, we need to ensure that the number of partitions created by the partitioner can easily divide the total number of devices."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HK2SfW9RoQjE"
      },
      "outputs": [],
      "source": [
        "if partitioner._model_parallel_submesh:\n",
        "  num_partitions = np.prod(partitioner._model_parallel_submesh)\n",
        "else:\n",
        "  num_partitions = partitioner._num_partitions\n",
        "if jax.device_count() % num_partitions != 0:\n",
        "  raise ValueError(\n",
        "    \"The number of devices available must be a multiple of the number of\",\n",
        "    f\" partitions. There are {jax.device_count()} devices available, but\",\n",
        "    f\" the number of partitions is set to {num_partitions}. Please\",\n",
        "    \" provide a different number of partitions.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C0sjeDUnoYz7"
      },
      "source": [
        "**Create a Checkpoint Manager**\n",
        "\n",
        "We make use of the T5X [`LegacyCheckpointManager`](https://github.com/google-research/t5x/blob/main/t5x/utils.py) to restore our model and save any future checkpoints. The `LegacyCheckpointManager` requires several constructor arguments:\n",
        "\n",
        "\n",
        "\n",
        "1.   `save_checkpoint_cfg`: an instance of the `SaveCheckpointConfig` wrapper class, which contains information about where and how to save future checkpoints.\n",
        "2.   `restore_checkpoint_cfg`: an instance of the `RestoreCheckpointConfig` wrapper class, which contains information and where and how to load checkpoints and restore the model.\n",
        "3.   `train_state_shape`: our model will load and save a T5X [`TrainState`](https://github.com/google-research/t5x/blob/main/t5x/train_state.py), which (as the name implies) stores information about the current state of training. We provide information about the shape of this train state to the checkpoint manager to enable saving this train state in checkpoints.\n",
        "4.   `partitioner`: our predefined partitioner.\n",
        "5.   `model_dir`: our previously configured output directory, where we will save any future checkpoints.\n",
        "\n",
        "Before we define these constructor arguments and initialize the checkpoint manager, let's discuss the T5X `TrainState` in a bit more depth. Importantly, T5X is a JAX-based library, which means that all of our methods follow typical functional programming patterns.\n",
        "\n",
        "\n",
        "Specifically, our training methods cannot have side effects, so we pass all model parameters, step number, optimizer state, etc. as input and get updated values as output from our methods. We use the T5X `TrainState` to hold all our model parameters, step number, optimizer state, etc. and we will later define a `train_step` method that will take in the train state and return an updated train state with new values.\n",
        "\n",
        "We define these constructor arguments and initialize the checkpoint manager below.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nC18r9Uaq2wL"
      },
      "outputs": [],
      "source": [
        "# Define CheckpointCfg wrappers.\n",
        "save_checkpoint_cfg = utils.SaveCheckpointConfig(\n",
        "        dtype=dtype,\n",
        "        keep=5, # The number of checkpoints to keep in the output_dir.\n",
        "        save_dataset=False)\n",
        "restore_checkpoint_cfg = utils.RestoreCheckpointConfig(\n",
        "        dtype=dtype,\n",
        "        mode=restore_mode,\n",
        "        path=checkpoint_path)\n",
        "\n",
        "# Define a train state initializer, which will help us get information about the\n",
        "# TrainState shape.\n",
        "train_state_initializer = utils.TrainStateInitializer(\n",
        "        optimizer_def=model.optimizer_def,\n",
        "        init_fn=model.get_initial_variables,\n",
        "        input_shapes=input_shapes,\n",
        "        input_types=None,\n",
        "        partitioner=partitioner)\n",
        "\n",
        "checkpoint_manager = utils.LegacyCheckpointManager(\n",
        "        save_cfg=save_checkpoint_cfg,\n",
        "        restore_cfg=restore_checkpoint_cfg,\n",
        "        train_state_shape=train_state_initializer.global_train_state_shape,\n",
        "        partitioner=partitioner,\n",
        "        ds_iter=None,\n",
        "        model_dir=output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSSfevHCsWX6"
      },
      "source": [
        "**Restore the Model from a Checkpoint or Initialize from Scratch** \\\n",
        "\n",
        "We try two different strategies for restoring a model. First, we try to restore the model from a checkpoint using the `CheckpointManager`. If no checkpoint can be found (likely because no path was provided in `checkpoint_path`), then we will initialize the model from scratch.\n",
        "\n",
        "Finally, we will log model initialization information (such as parameter shape, partitioning annotation, etc.) to the output directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-PgHNNOfuuFj"
      },
      "outputs": [],
      "source": [
        "def get_state(rng):\n",
        "  return train_state_initializer.from_scratch(rng).state_dict()\n",
        "\n",
        "# 1. Try to restore a model from a checkpoint.\n",
        "train_state = checkpoint_manager.restore(\n",
        "  [restore_checkpoint_cfg.path],\n",
        "  restore_checkpoint_cfg,\n",
        "  utils.get_fallback_state(restore_checkpoint_cfg, get_state, init_rng)\n",
        ")\n",
        "\n",
        "# 2. If no checkpoint to restore, init from scratch.\n",
        "if train_state is None:\n",
        "  train_state = train_state_initializer.from_scratch(init_rng)\n",
        "\n",
        "# Validate that we got an expected form of TrainState.\n",
        "if isinstance(train_state, Sequence):\n",
        "  raise ValueError(\n",
        "    \"Expected a single train state, but instead received a Sequence.\")\n",
        "train_state_axes = train_state_initializer.train_state_axes\n",
        "\n",
        "# Log the variable shapes information and write to a file.\n",
        "log_file = os.path.join(output_dir, \"model-info.txt\")\n",
        "utils.log_model_info(log_file,\n",
        "                     train_state_initializer.global_train_state_shape,\n",
        "                     partitioner)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JL5Nh6RvVbC"
      },
      "source": [
        "**Create a Trainer**\n",
        "\n",
        "Finally, we use many of the parameters we've defined above to create an instance of the T5X [Trainer](https://github.com/google-research/t5x/blob/main/t5x/trainer.py). The trainer takes in several constructor arguments:\n",
        "\n",
        "\n",
        "\n",
        "1.   `model`: the model that will be trained\n",
        "2.   `train_state`: a train state with parameters and optimizer state, which we've restored or initialized above.\n",
        "3.   `partitioner`: the partitioner to use.\n",
        "4.   `summary_dir`: the output directory, where we can write summaries of training.\n",
        "5.   `train_state_axes`: partitioning information for the optimizer, which we've initialized above.\n",
        "6.   `rng`: the JAX RNG to be used for training.\n",
        "7.   `learning_rate_fn`: a function that returns the learning rate given the current step. T5X provides some helper functions that define common learning rate schedules; we will use one of these helpers to define the learning rate in our example.\n",
        "\n",
        "We initialize a sample Trainer below.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pmVqBHrny2p3"
      },
      "outputs": [],
      "source": [
        "trainer = trainer_lib.Trainer(\n",
        "  model=model,\n",
        "  train_state=train_state,\n",
        "  partitioner=partitioner,\n",
        "  eval_names=[],\n",
        "  summary_dir=output_dir,\n",
        "  train_state_axes=train_state_axes,\n",
        "  rng=trainer_rng,\n",
        "  learning_rate_fn=utils.create_learning_rate_scheduler(),\n",
        "  num_microbatches=None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5_06UyFrlXFT"
      },
      "source": [
        "The code snippets above exactly replicate the `InteractiveModel` `__init__()` method (see [source code](https://github.com/google-research/t5x/blob/main/t5x/interactive_model.py)); running the code snippets above is exactly equivalent to running the single code snippet below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q_VA6DUi0MFv"
      },
      "outputs": [],
      "source": [
        "interactive_model = InteractiveModel(\n",
        "  batch_size=batch_size,\n",
        "  task_feature_lengths=task_feature_lengths,\n",
        "  output_dir=output_dir,\n",
        "  partitioner=partitioner,\n",
        "  model=model,\n",
        "  dtype=dtype,\n",
        "  restore_mode=restore_mode,\n",
        "  checkpoint_path=checkpoint_path,\n",
        "  input_shapes=input_shapes\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ANqpfv0lAVqL"
      },
      "source": [
        "**Defining a Batch of Examples to Train On**\\\n",
        "We are now ready to begin training!\n",
        "\n",
        "First, we'll begin by defining a batch of examples to train on; these examples should either be a list of inputs, or a list of dictionaries mapping 'target'/'input' keys to corresponding values, as shown below. For this Colab, we'll use a set of natural test questions and answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yhhR0yDcAn7w"
      },
      "outputs": [],
      "source": [
        "examples = [\n",
        "  {\n",
        "      'target': b'Ajay Tyagi',\n",
        "      'input':b'nq question: who has been appointed as the new chairman of sebi'\n",
        "  },\n",
        "  {\n",
        "      'target': b'C. S. Lewis',\n",
        "      'input': b'nq question: who wrote the book lion the witch and the wardrobe'},\n",
        "  {\n",
        "      'target': b'29',\n",
        "      'input': b'nq question: how many planes did japan lose at pearl harbor'},\n",
        "  {\n",
        "      'target': b'Jack Keil',\n",
        "      'input': b'nq question: who does the voice of mcgruff the dog'},\n",
        "  {\n",
        "      'target': b'Journey',\n",
        "      'input': b'nq question: who sings the wheels in the sky keep on turning'},\n",
        "  {\n",
        "      'target': b'Kumiko Watanabe',\n",
        "      'input': b'nq question: who voices regina in glitter force doki doki'},\n",
        "  {\n",
        "      'target': b'during World War II',\n",
        "      'input': b'nq question: when did the us become allies with britain'},\n",
        "  {\n",
        "      'target': b'the United States',\n",
        "      'input': b'nq question: who won the rugby 7 in las vegas'},\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jLc-P9S6_uOD"
      },
      "source": [
        "We also define the required features of the examples. For this Colab, we will only require an `inputs` and `targets` entry, as defined below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i5n6bnSq_2jF"
      },
      "outputs": [],
      "source": [
        "output_features = {\n",
        "        \"inputs\":\n",
        "            seqio.Feature(\n",
        "                vocabulary=model.input_vocabulary, add_eos=True),\n",
        "        \"targets\":\n",
        "            seqio.Feature(\n",
        "                vocabulary=model.output_vocabulary, add_eos=True)\n",
        "    }\n",
        "features = dict(sorted(output_features.items()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mixLzcBkQOT_"
      },
      "source": [
        "Now, let's (similarly) break down what the interactive model does when it takes a single step of training.\n",
        "\n",
        "The `InteractiveModel` `train_step()` method only performs two actions:\n",
        "\n",
        "\n",
        "1.   Convert the natural text examples into a tf.Dataset.\n",
        "2.   Take a single training step, using the T5X Trainer.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ug0zJx2kQk6g"
      },
      "source": [
        "**Prepare the dataset** \\\n",
        "\n",
        "Preparing the data for training is fairly straightforward. First, we validate that enough examples have been provided to train on a full batch of data.\n",
        "\n",
        "Then, we convert the natural text examples into a tf.Dataset and run any preprocessors; T5X has a helper function, `get_dataset_from_natural_text_examples`, that can do exactly that. For this example, the only preprocessing we will do is tokenization and appending an EOS token. If you are interested in learning more about preprocessors, please take a look at the [first Colab](https://colab.research.google.com/github/google-research/t5x/blob/main/t5x/notebooks/introduction.ipynb) in this tutorial series.\n",
        "\n",
        "Finally, we  convert all features using the model's feature converter and pad all batches of data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "chPomDFxQ6r3"
      },
      "outputs": [],
      "source": [
        "# Validate num examples.\n",
        "if len(examples) \u003c batch_size:\n",
        "  raise ValueError(\n",
        "    \"At least one batch of data must be provided. Please decrease the \"\n",
        "    \"batch_size or provide more examples.\")\n",
        "# Get a tf.Dataset.\n",
        "train_dataset = get_dataset_from_natural_text_examples(\n",
        "    examples=examples,\n",
        "    preprocessors=[\n",
        "        seqio.preprocessors.tokenize,\n",
        "        seqio.preprocessors.append_eos\n",
        "    ],\n",
        "    task_feature_lengths=task_feature_lengths,\n",
        "    features=features)\n",
        "\n",
        "# Convert and pad features.\n",
        "feature_converter = model.FEATURE_CONVERTER_CLS(pack=False)\n",
        "train_dataset = feature_converter(\n",
        "        train_dataset, task_feature_lengths=task_feature_lengths)\n",
        "train_dataset = train_dataset.padded_batch(batch_size, drop_remainder=True)\n",
        "train_iter = clu.data.dataset_iterator.TfDatasetIterator(train_dataset, checkpoint=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toq2uz7dAfdL"
      },
      "source": [
        "**Run 1 Training Step** \\\n",
        "\n",
        "We'll define a helper function that takes a single train step, making it easy to loop over this helper and train for multiple steps.\n",
        "\n",
        "Training is made fairly straightforward because of the T5X trainer. We'll simply add some logic to validate that it's ok for training to occur and to save a checkpoint. In total, we'll perform the following actions:\n",
        "\n",
        "\n",
        "1.   Validate that training can occur.\n",
        "2.   Take a training step.\n",
        "3.   Save a checkpoint.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qv8RXHeXK6mk"
      },
      "outputs": [],
      "source": [
        "def train_step(\n",
        "    trainer: t5x.trainer.Trainer,\n",
        "    train_state: t5x.train_state.TrainState,\n",
        "    train_iter: clu.data.dataset_iterator.TfDatasetIterator,\n",
        "    checkpoint_manager: utils.LegacyCheckpointManager,\n",
        "    save_checkpoint_cfg: utils.SaveCheckpointConfig):\n",
        "  # Validate that training can occur.\n",
        "  if trainer.stop_training:\n",
        "    logging.info(\"Stopping training early since `stop_training` is requested.\")\n",
        "    return\n",
        "\n",
        "  # Take a training step.\n",
        "  try:\n",
        "    first_step = int(utils.get_local_data(train_state.step))\n",
        "    train_summary = trainer.train(\n",
        "      train_iter, 1, start_step=first_step)\n",
        "  except trainer_lib.PreemptionError as e:\n",
        "    logging.info(\"Saving emergency checkpoint.\")\n",
        "    checkpoint_manager.save(\n",
        "      trainer.train_state,\n",
        "      save_checkpoint_cfg.state_transformation_fns)\n",
        "    logging.info(\"Saving emergency checkpoint done.\")\n",
        "    raise e\n",
        "\n",
        "  # Save a checkpoint.\n",
        "  logging.info(\"Saving checkpoint.\")\n",
        "  checkpoint_manager.save(\n",
        "      trainer.train_state,\n",
        "      save_checkpoint_cfg.state_transformation_fns)\n",
        "\n",
        "  # Wait until computations are done before exiting\n",
        "  multihost_utils.sync_global_devices(\"complete\")\n",
        "  return trainer.train_state, train_summary.result()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "srgktjMKMSZt"
      },
      "outputs": [],
      "source": [
        "print(f\"Current Step: {train_state.step}\")\n",
        "train_state, train_summary = train_step(trainer, train_state, train_iter, checkpoint_manager, save_checkpoint_cfg)\n",
        "print(f\"Current Step: {train_state.step}\")\n",
        "print(f\"Summary of Training: {train_summary}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AB0U_kfRNIyR"
      },
      "source": [
        "The code snippets above exactly replicate the `InteractiveModel` `train_step()` method (see [source code](https://github.com/google-research/t5x/blob/main/t5x/interactive_model.py)); running the code snippets above is exactly equivalent to running `interactive_model.train_step(examples)`.\n",
        "\n",
        "Alternately, you can loop over this helper function multiple times to finetune or pretrain a model (the code snippet below may take ~5 mins to run)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_7GJbHgNNRtl"
      },
      "outputs": [],
      "source": [
        "num_steps = 100\n",
        "for _ in range(num_steps):\n",
        "  # Reset the iterator, since we use the same batch for every step.\n",
        "  train_iter = clu.data.dataset_iterator.TfDatasetIterator(train_dataset, checkpoint=False)\n",
        "  train_state, train_summary = train_step(\n",
        "      trainer,\n",
        "      train_state,\n",
        "      train_iter,\n",
        "      checkpoint_manager,\n",
        "      save_checkpoint_cfg\n",
        "  )\n",
        "print(f\"Current Step: {train_state.step}\")\n",
        "print(f\"Summary of Training: {train_summary}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gs3BwU5sKRmu"
      },
      "source": [
        "The code snippets above demonstrate how T5X runs training. You can exactly replicate this behavior by using the `InteractiveModel`, as described above."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lcDwmp_AxnOG"
      },
      "source": [
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-QR5LnmN4ikp"
      },
      "source": [
        "# Advanced Topics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CLstCKpP8Ge7"
      },
      "source": [
        "## T5X Training Binaries and Other Advanced Features"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h9RT_9R_8K3U"
      },
      "source": [
        "T5X offers training binaries that have the same functionality as the InteractiveModel, with additional features as well (more advanced compiling, custom checkpointing periods, etc.). Importantly, these binaries are configured using [Gin](https://github.com/google/gin-config/blob/main/README.md); if you are not familiar with Gin, please take a look at this [Gin Primer](https://github.com/google-research/t5x/blob/main/docs/usage.md/gin) to get started.\n",
        "\n",
        "If you are familiar with Gin and interested in using the T5X training binaries, we have provided a helper function, `get_gin_config_from_interactive_model`, which will take an InteractiveModel instance and generate the gin config that you can use to run the T5X training binaries; this gin config will exactly reproduce the InteractiveModel training functionality we've described above. We've provided an example below.\n",
        "\n",
        "Importantly, the InteractiveModel takes in a model, partitioner, and data, so we cannot generate Gin configs for these components. You can pass Gin config strings for the model and partitioner components to the helper function, as demonstrated below. Additionally, you can pass a SeqIO task containing your data to the helper function. See the section below if you are unfamiliar with SeqIO."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rhgUZ0w6yQsE"
      },
      "outputs": [],
      "source": [
        "# Define an InteractiveModel instance, based on the `tiny` T5X EncoderDecoder model.\n",
        "input_shapes = {\n",
        "    'encoder_input_tokens': np.array([8, 38]),\n",
        "    'decoder_target_tokens': np.array([8, 18]),\n",
        "    'decoder_input_tokens': np.array([8, 18]),\n",
        "    'decoder_loss_weights': np.array([8, 18])\n",
        "}\n",
        "t5_config = network.T5Config(\n",
        "    vocab_size=32128,\n",
        "    dtype='bfloat16',\n",
        "    emb_dim=8,\n",
        "    num_heads=4,\n",
        "    num_encoder_layers=2,\n",
        "    num_decoder_layers=2,\n",
        "    head_dim=3,\n",
        "    mlp_dim=16,\n",
        "    mlp_activations=('gelu', 'linear'),\n",
        "    dropout_rate=0.0,\n",
        "    logits_via_embedding=False)\n",
        "module = network.Transformer(config=t5_config)\n",
        "model = t5x.models.EncoderDecoderModel(\n",
        "    module=module,\n",
        "    input_vocabulary=t5.data.get_default_vocabulary(),\n",
        "    output_vocabulary=t5.data.get_default_vocabulary(),\n",
        "    optimizer_def=t5x.adafactor.Adafactor(decay_rate=0.8, step_offset=0),\n",
        "    decode_fn=functools.partial(\n",
        "        t5x.decoding.temperature_sample, temperature=1.0, topk=40))\n",
        "interactive_model = InteractiveModel(\n",
        "      batch_size=8,\n",
        "      task_feature_lengths={\n",
        "            'inputs': 32,\n",
        "            'targets': 32\n",
        "        },\n",
        "      output_dir='/tmp',\n",
        "      partitioner=partitioning.PjitPartitioner(\n",
        "        num_partitions=2,\n",
        "        model_parallel_submesh=None,\n",
        "        logical_axis_rules=partitioning.standard_logical_axis_rules()),\n",
        "      model=model,\n",
        "      dtype='float32',\n",
        "      restore_mode='specific',\n",
        "      checkpoint_path='',\n",
        "      input_shapes=input_shapes,\n",
        "      input_types=None)\n",
        "\n",
        "# Define Gin Config strings for the model, partitioner, and any imports.\n",
        "imports_str = \"\"\"from t5x import models\n",
        "from t5x import partitioning\n",
        "import t5.data.mixtures\n",
        "include 't5x/examples/t5/t5_1_1/tiny.gin'\"\"\"\n",
        "partitioner_config = 'partitioning.PjitPartitioner.num_partitions = 2'\n",
        "model_config = \"\"\"models.EncoderDecoderModel:\n",
        "  z_loss = 0.0\n",
        "  label_smoothing = 0.0\n",
        "  loss_normalizing_factor = None\"\"\"\n",
        "\n",
        "gin_config_str = get_gin_config_from_interactive_model(\n",
        "    interactive_model=interactive_model,\n",
        "    script_type=T5XScriptType.PRETRAINING,\n",
        "    task_name='wmt19_ende_v003',\n",
        "    partitioner_config_str=partitioner_config,\n",
        "    model_config_str=model_config,\n",
        "    train_steps=3,\n",
        "    imports_str=imports_str)\n",
        "print(gin_config_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uGd1DxDT3gB7"
      },
      "source": [
        "\n",
        "Once you have generated the `gin_config_str` as above, you can write this string to a file and launch your training experiment locally by running the following on commandline:\n",
        "\n",
        "\n",
        "```\n",
        "MODEL_DIR=\"/tmp/pretrain-model/\"\n",
        "python -m t5x.train \\\n",
        "  --gin_file=${GIN_FILE_PATH} \\\n",
        "  --gin.MODEL_DIR=\\\"${MODEL_DIR}\\\" \\\n",
        "  --alsologtostderr\n",
        "```\n",
        "\n",
        "For more details on training using the T5X training binaries, please see the [Pretraining](https://github.com/google-research/t5x/blob/main/docs/usage.md/pretrain) or [Finetuning](https://github.com/google-research/t5x/blob/main/docs/usage.md/finetune) tutorials."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wi29fMdv4mSr"
      },
      "source": [
        "## SeqIO\n",
        "\n",
        "If you are interested in T5X, you may also be interested in, or have heard of, SeqIO. SeqIO is a library for processing sequential data to be fed into downstream sequence models. At a high level, SeqIO relies on user-defined `Tasks` and `Mixtures` that can be used to retrieve and evaluate datasets.\n",
        "\n",
        "We won't go into details about SeqIO here; we recommend checking out this [SeqIO Introductory guide](https://github.com/google/seqio/blob/main/README.md/index) and/or clicking below to run a SeqIO Introductory Colab. The rest of this section will assume a basic understanding of SeqIO.\n",
        "\n",
        "\u003ca href=\"https://colab.research.google.com/github/google-research/seqio/blob/main/seqio/notebooks/Basics_Task_and_Mixtures.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n",
        "\n",
        "If you are already familiar with SeqIO and have a SeqIO task/mixture that you would like to use in this Colab, we do provide a SeqIO bridge that takes in a SeqIO task/mixture and produces batches of examples that can be processed by the code snippets above. We've provided an example of this bridge below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kh5VMAWcV_WG"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/google-research/google-research.git google_research"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bM0nRIEFwyj_"
      },
      "outputs": [],
      "source": [
        "import google_research.t5_closed_book_qa.t5_cbqa.tasks\n",
        "batches = get_batches_from_seqio(\n",
        "        task_or_mixture_name='natural_questions_open',\n",
        "        split='validation',\n",
        "        batch_size=8,\n",
        "        num_batches=2,\n",
        "        seed=42)\n",
        "print(f\"Batches: {batches}\")\n",
        "# Train the interactive model on the provided batches.\n",
        "original_step = interactive_model.step\n",
        "_ = interactive_model.train_loop(num_steps=len(batches), train_batches=batches)\n",
        "print(f\"Original Step: {original_step}, Current Step: {interactive_model.step}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Elt08160w03X"
      },
      "source": [
        "The `get_batches_from_seqio` bridge can take several constructor arguments:\n",
        "\n",
        "\n",
        "1.   `task_or_mixture_name`: the name of the SeqIO task/mixture to read data from. It should be noted that your task/mixture must already be registered with SeqIO, and you must import the module that defines your task/mixture here (as seen above).\n",
        "2.   `split`: the split of the Task/Mixture to read data from.\n",
        "3.   `batch_size`: how many examples should appear in each batch.\n",
        "4.   `num_batches`: the total number of batches to return.\n",
        "5.   `get_pretokenized_examples`: optional. A boolean, defaulting to True, that determines whether we should read the `inputs_pretokenized`/`targets_pretokenized` elements from an example, or the `inputs`/`targets` elements. \\\n",
        "The `train_step`, `predict`, `predict_with_aux`, `score`, and `evaluate` methods of the InteractiveModel assume that we should run [tokenization](https://github.com/google/seqio/tree/main/seqio/preprocessors.py) and [appending an EOS token](https://github.com/google/seqio/tree/main/seqio/preprocessors.py) as the only preprocessors. To use these methods with this pre-defined list of preprocessors, you can set `get_pretokenized_examples=True` to retrieve examples that still need to be tokenized, and these InteractiveModel methods will handle running these preprocessors. This setting can also be helpful if you want to inspect the natural text inputs/targets of your SeqIO task. \\\n",
        "However, some SeqIO tasks do not use tokenization (ex: span corruption). You can set `get_pretokenized_examples=False`, and this bridge will read the fully preprocessed examples from the SeqIO task. You can then run `train_step_with_preprocessors`, `infer_with_preprocessors`, or `evaluate_with_preprocessors` and provide an empty preprocessors list (because all preprocessing has already been completed by this bridge) to run training/inference/evaluation. We have provided an example of using this bridge to retrieve fully preprocessed examples below.\n",
        "6.   `sequence_length`: optional. A dictionary mapping feature key to maximum length (int) for that feature. Used by SeqIO to retrieve the dataset/examples.\n",
        "7.   `**get_dataset_kwargs`: there are many [additional parameters](https://github.com/google/seqio/tree/main/seqio/dataset_providers.py) that can be set in the `SeqIO.get_dataset` function. If you would like to set any of these arguments, you can set them using this `kwargs` parameter.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fjKBCX39w0Xl"
      },
      "outputs": [],
      "source": [
        "import t5.data.tasks\n",
        "batches = get_batches_from_seqio(\n",
        "    task_or_mixture_name='c4_v220_span_corruption',\n",
        "    split='validation',\n",
        "    batch_size=8,\n",
        "    num_batches=1,\n",
        "    get_pretokenized_examples=False,\n",
        "    sequence_length=interactive_model._task_feature_lengths,\n",
        "    seed=42)\n",
        "batch = batches[0]  # We expect only a single batch.\n",
        "original_step = interactive_model.step\n",
        "interactive_model.train_step_with_preprocessors(\n",
        "        examples=batch, preprocessors=[])\n",
        "print(f\"Original Step: {original_step}, Current Step: {interactive_model.step}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "Welcome to T5X: Training Deep Dive",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1hQO9MD6psZtTeqZyXPJIoUV0uzTa2qPg",
          "timestamp": 1662951508591
        },
        {
          "file_id": "1Akpc6pKlJB5rn5YYYFC9lw2OMk6oBzlQ",
          "timestamp": 1662754223629
        },
        {
          "file_id": "1rA8bgO2bJRoebAuS96Ji0RUhnawgBY4i",
          "timestamp": 1650477076639
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
